import pygame
from .constants import *
from game.game_mode import GameMode, RobotDifficulty


class Menu:
    def __init__(self, screen, font):
        self.screen = screen
        self.font = font
        self.large_font = pygame.font.Font(None, 48) if font is None else font

        # 菜单状态
        self.current_menu = "main"  # main, mode_select, difficulty_select

        # 按钮定义
        button_width = 200
        button_height = 50
        button_spacing = 20
        start_y = WINDOW_HEIGHT // 2 - 50
        center_x = WINDOW_WIDTH // 2

        # 主菜单按钮
        self.pvp_button = pygame.Rect(center_x - button_width // 2, start_y, button_width, button_height)
        self.pvc_button = pygame.Rect(center_x - button_width // 2, start_y + button_height + button_spacing, button_width, button_height)

        # 难度选择按钮
        self.easy_button = pygame.Rect(center_x - button_width // 2, start_y - button_height - button_spacing, button_width, button_height)
        self.medium_button = pygame.Rect(center_x - button_width // 2, start_y, button_width, button_height)
        self.openai_button = pygame.Rect(center_x - button_width // 2, start_y + button_height + button_spacing, button_width, button_height)

        # 返回按钮
        self.back_button = pygame.Rect(50, WINDOW_HEIGHT - 80, 100, 40)

    def handle_click(self, pos):
        """处理菜单点击事件，返回选择的游戏模式和难度"""
        if self.current_menu == "main":
            if self.pvp_button.collidepoint(pos):
                return GameMode.PVP, None
            elif self.pvc_button.collidepoint(pos):
                self.current_menu = "difficulty_select"
                return None, None

        elif self.current_menu == "difficulty_select":
            if self.easy_button.collidepoint(pos):
                return GameMode.PVC, RobotDifficulty.EASY
            elif self.medium_button.collidepoint(pos):
                return GameMode.PVC, RobotDifficulty.MEDIUM
            elif self.openai_button.collidepoint(pos):
                return GameMode.PVC, RobotDifficulty.OPENAI
            elif self.back_button.collidepoint(pos):
                self.current_menu = "main"
                return None, None

        return None, None

    def draw(self, mouse_pos):
        """绘制菜单界面"""
        self.screen.fill(BOARD_COLOR)

        # 绘制标题
        title_text = self.large_font.render("五子棋", True, TEXT_COLOR)
        title_rect = title_text.get_rect(center=(WINDOW_WIDTH // 2, 150))
        self.screen.blit(title_text, title_rect)

        if self.current_menu == "main":
            self._draw_main_menu(mouse_pos)
        elif self.current_menu == "difficulty_select":
            self._draw_difficulty_menu(mouse_pos)

    def _draw_main_menu(self, mouse_pos):
        """绘制主菜单"""
        self._draw_button(self.pvp_button, "人人对战", mouse_pos)
        self._draw_button(self.pvc_button, "人机对战", mouse_pos)

    def _draw_difficulty_menu(self, mouse_pos):
        """绘制难度选择菜单"""
        # 副标题
        subtitle_text = self.font.render("选择机器人难度", True, TEXT_COLOR)
        subtitle_rect = subtitle_text.get_rect(center=(WINDOW_WIDTH // 2, 250))
        self.screen.blit(subtitle_text, subtitle_rect)

        # 难度按钮
        self._draw_button(self.easy_button, "简单", mouse_pos)
        self._draw_button(self.medium_button, "中等", mouse_pos)
        self._draw_button(self.openai_button, "AI大模型", mouse_pos)

        # 返回按钮
        self._draw_button(self.back_button, "返回", mouse_pos)

    def _draw_button(self, rect, text, mouse_pos):
        """绘制单个按钮"""
        color = BUTTON_HOVER if rect.collidepoint(mouse_pos) else BUTTON_COLOR
        pygame.draw.rect(self.screen, color, rect, border_radius=8)

        # 绘制边框
        pygame.draw.rect(self.screen, LINE_COLOR, rect, 2, border_radius=8)

        # 绘制文字
        text_surface = self.font.render(text, True, TEXT_COLOR)
        text_rect = text_surface.get_rect(center=rect.center)
        self.screen.blit(text_surface, text_rect)